#include "./nn_yolt/yolt_def.hpp"
#include "./nn_yolt/yolt.h"
#include "./tools/generate_yolt_train_data.hpp"

#include <ctime>
#include <fstream>
#include <vector>
#include <iostream>

#define TEST false

static const std::string log_file_name = "yolt.log";

static void save_log(int curr_train_index)
{
    // 获取当前时间
    time_t curr_time;
    std::time(&curr_time);
    char tmp[64];
    std::strftime(tmp, sizeof(tmp), "%Y-%m-%d %H:%M:%S", std::localtime(&curr_time));

    std::cout << tmp << std::endl;

    std::ofstream f_out;
    f_out.open(log_file_name.c_str(), std::ios_base::app | std::ios_base::out);

    assert(f_out.is_open());

    f_out << "time: " << tmp << std::endl;
    f_out << "train index: " << curr_train_index << std::endl;

    f_out.close();
}

static int load_log()
{
    std::ifstream f_in;
    f_in.open(log_file_name, std::ios_base::ate);

    if (!f_in.is_open())
    {
        throw std::exception();
    }

    f_in.seekg(-2, f_in.cur);

    // 搜索倒数第二行的换行符
    while (f_in.peek() != f_in.widen('\n'))
    {
        f_in.seekg(-1, f_in.cur);
    }
    // 走到这里表示跳过一行了，所以前进两个字符，跳过该换行符，准备读取数据
    f_in.seekg(2, f_in.cur);
    std::string line;
    std::getline(f_in, line);
    f_in.close();

    std::vector<int> nums;
    std::vector<std::string> strs;
    tools::extract_info_from_str_line(line, nums, strs, " ");

    return nums[0];
}

#if TEST

int main(const int argc, const char *argv[])
{

#if GENERATE_YOLT_TRAIN_DATA_TEST
    tools::TestGenerateYoltTrainData::run();
#endif

    return 0;
}

#endif

#if !TEST
int main(const int argc, const char *argv[])
{
    std::cout << "NNYolt" << std::endl;
    if (argc <= 1)
    {
        return -1;
    }

    std::string opt = argv[1];

    const int CELL_NX = 13;
    const int CELL_NY = 13;
    const int BOX_NUM = 4;
    const int IMAGE_WIDTH = 416;
    const int IMAGE_HEIGHT = 416;
    const int IMAGE_CHANNELS = 3;
    const std::string net_file_name = "net.dat";
    const std::string TRAIN_DATA_ROOT_DIR = "../../../datas/train/";
    const std::string IMAGES_DIR = TRAIN_DATA_ROOT_DIR + "images/";
    const std::string LABEL_DIR = IMAGES_DIR + "label/";

    nn_yolt::NNYolt *yolt = new nn_yolt::NNYolt(IMAGE_WIDTH, IMAGE_HEIGHT, IMAGE_CHANNELS,
                                                CELL_NX, CELL_NY, BOX_NUM, net_file_name);

    if (opt == "gen")
    {
        tools::GenerateYoltTrainData train_data_generator({IMAGE_WIDTH, IMAGE_HEIGHT, IMAGE_CHANNELS},
                                                          CELL_NX, CELL_NY, BOX_NUM, 0.15);
        train_data_generator.generate(TRAIN_DATA_ROOT_DIR);
    }
    else if (opt == "train")
    {
        const std::string train_file_path = TRAIN_DATA_ROOT_DIR + "yolt";

        int batch_size = 1;
        int epoch = 1;
        int train_num = 10;
        int train_index_begin = 0;

        if (argc >= 3)
        {
            batch_size = std::stoi(argv[2]);
        }
        if (argc >= 4)
        {
            epoch = std::stoi(argv[3]);
        }
        if (argc >= 5)
        {
            train_index_begin = std::stoi(argv[4]);
        }
        if (argc >= 6)
        {
            train_num = std::stoi(argv[5]);
        }

        try
        {
            yolt->load_net(net_file_name);
            std::cout << "Load net success!" << std::endl;
        }
        catch (const std::exception &e)
        {
            yolt->construct_net();
            std::cout << "Construct net success!" << std::endl;
            yolt->save_net(net_file_name);
        }

        yolt->train(train_file_path, train_index_begin, train_num, batch_size, epoch);
    }
    else if (opt == "test")
    {
        const std::string test_file_path = TRAIN_DATA_ROOT_DIR + "yolt";
        int test_index_begin = 0;
        int test_num = 10;
        if (argc >= 3)
        {
            test_index_begin = std::stoi(argv[2]);
        }
        if (argc >= 4)
        {
            test_num = std::stoi(argv[3]);
        }
        try
        {
            yolt->load_net(net_file_name);
            std::cout << "Load net success!" << std::endl;
            yolt->test(test_file_path, test_index_begin, test_num);
        }
        catch (const std::exception &e)
        {
            std::cout << "Load net failed!" << std::endl;
            return 0;
        }
    }
    else if (opt == "help")
    {
        std::cout << "train -batch_size -epoch -train_begin_index -train_num" << std::endl;
        std::cout << "test -test_begin_index -test_num" << std::endl;
    }

    delete yolt;

    return 0;
}
#endif // #if !TEST